/*12/14/09 20:23
*
*Cha Li
*CS4804 - Final Project
*Developed in Ubuntu 9.04 x64
*Compiled using g++
*
*/

#include <cstdlib>
#include <iostream>
#include <fstream>
#include <cmath>
#include <list>
#include "Square.h"
#include "EquationHelper.h"
using namespace std;

#define GAMMA .2	//the discount rate for future rewards
#define ALPHA .1	//used for temporal difference
#define OUT_TYPE results //cout or results
#define EDGE_IS_PLACED(x,y)	(x & (int)pow(2,y))

ofstream results("output",ios::trunc);

//all the tables
float **q_table;
float **pi_table;
float **rw_table;

//game variables 
int TRAIN_GAMES;
int NUM_GAMES;
char mode;
Square *board;
int rows, cols;
int total_states;
int total_edges;
int total_squares;
int turn;
int h_start;

//loop variables
int cur_edge, cur_state;

//function prototypes
void initRwTable(void);
void initQTable(void);
void initPiTable(void);

//fill in rows as they are encountered and needed
//SAVES MEMORY!!!
void fillRwTableState(int, int);
void fillQTableState(int,int);
void fillPiTableState(int,int);

void updateQTable(int);
void updatePiTable(void);

void printTable(float **);
int compTurn(int*);

void temporal_check(void);
int edgeToSquare(int, int);
int getPoints(int, int);

//a return list item
struct reward_item{
	int state;
	int action;
};

//return list
list<reward_item> reward_list;

bool training = true;

int main(int argc, char **argv) {
	srand(time(0));
	if(argc < 5){
		cout << "Incorrect Args\nSyntax: driver [rows] [cols] [num. training games][num. test games][options]  " << endl;
		cout << "options:\n\t t\tUse temporal difference algorithm";
	}	
	else{
		rows = atoi(argv[1]); //m
		cols = atoi(argv[2]); //n
		TRAIN_GAMES = atoi(argv[3]); //number of training games
		NUM_GAMES = atoi(argv[4]);
		//check for temporal difference
		if(argc == 6){
			if(*argv[5] == 't')
				mode = *argv[5];
			else
				cout << "Invalid option: " << *argv[5] << endl;
		}

		//board stats
		total_squares = rows * cols;
		int board_state = 0;	
		total_edges = (rows + 1) * cols + (cols + 1) * rows;
		total_states = pow(2, total_edges);
		float win_percentage = 0;
	
		cout << "Initializing Game...\n";
		if(mode == 't')
			cout << "*Temporal Difference Enabled\n";
		results << "Total Squares: " << total_squares << " Total Edges: " << total_edges << " Total States :  " << total_states << endl;
		cout << "Total Squares: " << total_squares << " Total Edges: " << total_edges << " Total States :  " << total_states << endl;

		//index of the first horizontal edge
		h_start	= rows * (cols + 1);
		
		//create the board	
		OUT_TYPE << "Creating Board...\n";
		board = new Square[total_squares];
		int cur_square;
		for(cur_square = 0; cur_square < total_squares; cur_square++)
			board[cur_square].initSquare(cur_square, rows, cols, h_start);

		//initialize all the tables
		initRwTable();
		initQTable();
		initPiTable();

		int your_move, your_score;
		int cpu_move, cpu_score, cpu_wins = 0, cpu_losses = 0;;
		int t_game = 0; int n_game;
		OUT_TYPE << "Training the computer...\n\n";
		OUT_TYPE << "t_game\tGames Played\tWins\tLosses\tTies\n";
		while(t_game < TRAIN_GAMES){
			board_state = 0;
			your_score = 0; cpu_score = 0;
			turn = rand() % 2; //see who goes first
			while(board_state != (total_states - 1)){
				if(turn == 0){
					your_move = rand() % total_edges;				//computer trains against a random edge placer
					if(!EDGE_IS_PLACED(board_state, your_move)){ //make sure edge is valid
						board_state += (int)pow(2, your_move);
						if( getPoints(your_move, board_state) == 0) //no score, no extra move
							turn = 1;
						your_score += getPoints(your_move, board_state);
					}
					else
						continue;
				}
				else{
					if(board_state == (total_states - 1)) //game is over
						break;
					cpu_move = compTurn(&board_state);
					if( getPoints(cpu_move, board_state) == 0)	//no score, no extra move
						turn = 0;
					cpu_score += getPoints(cpu_move, board_state);
				}
			}
			//computer 'learns' at this point
			if(cpu_score > your_score)
				updateQTable(300);		//cpu won, it gets a bonus reward of 300
			else if(cpu_score < your_score)
				updateQTable(0);		//cpu lost, it gets no bonus
			else
				updateQTable(50);		//tie game, still gets a small bonus since it didn't lose
			updatePiTable(); 

			t_game++;
			n_game = 0;

			//test what the computer has learned//
			cpu_wins = 0; cpu_losses = 0;
			while(n_game < NUM_GAMES){
				your_move = rand() % total_edges;
				board_state = 0; your_score = 0; cpu_score = 0; int points;
				turn = rand() % 2;		//determine who goes first
				while(board_state != (total_states - 1)){

					if(turn == 0){
						your_move = rand() % total_edges;
						if(!EDGE_IS_PLACED(board_state, your_move)){
							board_state += (int)pow(2, your_move);
							points = getPoints(your_move, board_state);
							if( points == 0) //no points, no extra move
								turn = 1;
							your_score += points;
						}
						else
							continue;
					}
					else{	 //cpu's turn
						if(board_state == (total_states - 1))
							break;
						cpu_move = compTurn(&board_state);
						points = getPoints(cpu_move, board_state);
						if( points == 0)	//no points, no extra move
							turn = 0;
						cpu_score += points;
					}
				}

				//keep track of wins and losses
				if(cpu_score > your_score)
					cpu_wins++;
				else if(cpu_score < your_score)
					cpu_losses++;
				n_game++;
				reward_list.clear();
			}
			win_percentage += cpu_wins;
			OUT_TYPE << t_game << "\t" << NUM_GAMES << "\t\t" << cpu_wins << "\t" << cpu_losses << "\t" << (NUM_GAMES - cpu_wins - cpu_losses) << endl;
		}
		results << "G: " << GAMMA << " A: " << ALPHA << " Win Percentage: " << (win_percentage/(NUM_GAMES * TRAIN_GAMES)) << endl;
		cout << "G: " << GAMMA << " A: " << ALPHA << " Win Percentage: " << (win_percentage/(NUM_GAMES * TRAIN_GAMES)) << endl;
		cout << "\nDone. Please check file \"output\" for results.\n";

		/*NOTE: printTable is disabled if there are too many states. 2^22 is the cut off point */
		OUT_TYPE << "\n---Training Results----\nRewards Table:\n";
		printTable(rw_table);
		OUT_TYPE << "Q Table:\n";
		printTable(q_table);
		OUT_TYPE << "Pi Table:\n";
		printTable(pi_table);
	}
	return 0;
}
/******************* END MAIN ****************************/


/* only allocate memory for one dimension of each 2d array.
 * the rows will be filled in as the computer encounters the states.
 * DONT ALLOCATE MEMORY ALL AT ONCE, IF YOU DONT NEED TO USE ALL OF IT
 * SAVES MEMORY LIKE CRAAAZY
 */
void initRwTable(){
	rw_table = new float *[total_states];
}

void initQTable(){
	q_table = new float* [total_states];
}

void initPiTable(){
	pi_table = new float* [total_states];
}

void fillRwTableState(int cur_state){
	int sqr_i;	//index of a square
	rw_table[cur_state] = new float[total_edges];
	for(cur_edge = 0; cur_edge < total_edges; cur_edge++){
		int w1 = 0, w2 = 0; 	//square worths
		sqr_i = -1;
		//process vertical edges
		if(cur_edge < h_start){
			//only find the square to the left of an edge if the edge
			//is not the left most edge
			if(!isLeftEdge(cur_edge, cols)){
				sqr_i = edgeToSquare(sqr_i, cur_edge);
				w1 = board[sqr_i].getWorth(cur_state);
			}

			//only find the square to the right of an edge if the edge
			//is not the right most edge
			if(!isRightEdge(cur_edge, cols)){
				sqr_i = edgeToSquare(sqr_i, cur_edge);
				w2 = board[sqr_i].getWorth(cur_state);
			}
		}
		//process horizontal edges
		else{
			//only find the square above an edge if the edge is
			//not in the top most position
			if(!isTopEdge(cur_edge, rows, h_start)){
				sqr_i = edgeToSquare(sqr_i, cur_edge);
				w1 = board[sqr_i].getWorth(cur_state);
			}

			//only find the square below an endge if the edge is
			//	not in the bottom most position
			if(!isBottomEdge(cur_edge, rows, h_start)){
				sqr_i = edgeToSquare(sqr_i, cur_edge);
				w2 = board[sqr_i].getWorth(cur_state);
			}
		}

		//only put a value in the reward matrix if the chosen action
		//	doesnt lead back to the same state. (edge is already placed)
		if(!EDGE_IS_PLACED(cur_state, cur_edge)){
			//202 represents a move that sets of a chain square completions
			//	this is the best move to make if available
			if((w1 + w2) == 202)
				rw_table[cur_state][cur_edge] = 500;
			else
				rw_table[cur_state][cur_edge] = w1 + w2;
		}
		else
			rw_table[cur_state][cur_edge] = 0;
	}
}

//Initializes the Q-Values for a given state to random values
void fillQTableState(int cur_state){
	q_table[cur_state] = new float[total_edges];
	for(cur_edge = 0; cur_edge < total_edges; cur_edge++)
				q_table[cur_state][cur_edge] = (rand() % 20);
}

//Initializes the Pi Values for a given state using an epsilon
//soft policy of .9/.1
void fillPiTableState(int cur_state){
	if(cur_state != total_states - 1){
		int best; float low; int placed; float the_chosen;
		pi_table[cur_state] = new float[total_edges];
		do{
			best = rand() % total_edges; //pick the edge index that will get the .9
		}while(EDGE_IS_PLACED(cur_state, best));

		//first run-through, set everything to 0 and keep track of placed edges
		placed = 0;
		for(cur_edge = 0; cur_edge < total_edges; cur_edge++){
			pi_table[cur_state][cur_edge] = 0;
			if(EDGE_IS_PLACED(cur_state, cur_edge))
				placed++;
		}
		the_chosen = ((total_edges - 1) == placed) ? 1.0 : .9;
		low = .1 / (total_edges - placed - 1);
		for(cur_edge = 0; cur_edge < total_edges; cur_edge++)
			if(!EDGE_IS_PLACED(cur_state, cur_edge))
				pi_table[cur_state][cur_edge] = (cur_edge == best) ? the_chosen : low;
	}
	else{
		//handle last state (filled bored) all %'s = 0
		pi_table[total_states - 1] = new float[total_edges];
		for(cur_edge = 0; cur_edge < total_edges; cur_edge++)
			pi_table[total_states - 1][cur_edge] = 0;
	}
}
/****** END FILL TABLE FUNCTIONS *******/


/*determines what edge to place given the current state of the board. used by
	 the computer player */
int compTurn(int *state){
	float mv_percent = ((rand() % 100) / 100.0); //calculate a percent, used to determing which move is chosen  
	//mv_percent is 0, it must be bumped to a non-zero number. 
	//	otherwise action with 0% probability will get chosen
	mv_percent = (mv_percent == 0) ? .01 : mv_percent;

	int edge_chosen = -1; float prob_sum = 0;
	struct reward_item this_move;

	if(pi_table[*state] == NULL)	//initialize needed part of the table
		fillPiTableState(*state);

	for(cur_edge = 0; cur_edge < total_edges; cur_edge++){
		prob_sum += pi_table[*state][cur_edge];
		if(mv_percent <= prob_sum){
			edge_chosen = cur_edge;
			break;
		}
	}

	//the program should never go in here during normal execution.
	//if it does goes in here, we have a problem with the Q and/or PI
	//tables.
	if(edge_chosen == -1){
		cout << "Offending State: " << *state << endl;
		cout << "Skipped Turn. YOU SHOULD NOT BE SEEING THIS\n";
		cout << "\nState dumped.\n";
		/*NOTE: printTable is disabled if there are too many states. 2^22 is the cut off point */
		printTable(rw_table);
		printTable(q_table);
		printTable(pi_table);
		exit(0);
	}

	//create a item to be stored in the master list
	this_move.state = *state;
	this_move.action = edge_chosen;
	reward_list.push_front(this_move);
	// > .1 means only do the temporal_check if the chosen action was exploitation
	//anything <= .1 means it was done out of curiosity
	if(mode == 't' && pi_table[this_move.state][this_move.action] > .1){
		temporal_check();
	}

	//update board state
	*state += (int)pow(2, edge_chosen);
	return edge_chosen;
}


/*Q-value update algorithms*/
void temporal_check(){ 
	int actual_q = 0;
	list<reward_item>::iterator iter;
	reward_item first, cur;
	//if theres only item in the reward list, it will obviously be equal to itself. //
	if(reward_list.size() > 1){
	  first = reward_list.front();
	  reward_list.pop_front();
		if(rw_table[first.state] == NULL) //initialize needed part of the table
			fillRwTableState(first.state);
	 
		 actual_q = rw_table[first.state][first.action];
	  for( iter = reward_list.begin(); iter != reward_list.end(); ++iter) {
		  cur = *iter;
			if(rw_table[cur.state] == NULL) //initialize needed part of the table
				fillRwTableState(cur.state);
			
			if(q_table[cur.state] == NULL) //initialize needed part of the table
				fillQTableState(cur.state);
			//this is the crazy equation given in the pseudocode. as you can see, it's crazy.
			actual_q = q_table[cur.state][cur.action] + ALPHA * ((rw_table[cur.state][cur.action] + GAMMA * actual_q) - q_table[cur.state][cur.action]);

			if(q_table[cur.state] == NULL) //initialize needed part of the table
				fillQTableState(cur.state);
		  if(q_table[cur.state][cur.action] != actual_q)
			  q_table[cur.state][cur.action] = actual_q;
	  }
	  reward_list.push_front(first);
	}
}

/*depending on the outcome of a game, bonus can be a + value (won or tie)
	or 0 (lost) */
void updateQTable(int bonus){
	float total_reward = 0;

	/*handle the first item in the list (last action taken) seperately from the rest */
	reward_item cur = reward_list.front();
	if(rw_table[cur.state] == NULL) //initialize needed part of the table
		fillRwTableState(cur.state);
	total_reward = rw_table[cur.state][cur.action] + bonus;
	if(q_table[cur.state] == NULL) //initialize needed part of the table
		fillQTableState(cur.state);
	q_table[cur.state][cur.action] = total_reward;
	reward_list.pop_front();

	while(!reward_list.empty()){
		cur = reward_list.front();
		if(rw_table[cur.state] == NULL) //initialize needed part of the table
				fillRwTableState(cur.state);
		total_reward = rw_table[cur.state][cur.action] + (GAMMA * total_reward);
		if(q_table[cur.state] == NULL) //initialize needed part of the table
				fillQTableState(cur.state);
		q_table[cur.state][cur.action] = (q_table[cur.state][cur.action] + total_reward) / 2;
		reward_list.pop_front();
	}
}

/*redistribute the percentages, and get ready for the next game*/
void updatePiTable(){ 
	for(cur_state = 0; cur_state < total_states; cur_state++){
		int placed = 0; float low;
		int q_high = 0, q_count = 0, best = 0;
		float the_chosen;
		if(q_table[cur_state] == NULL) //initialize needed part of the table
			continue;
		else{
			/*find the maximum Q value, if there are multiple instances of the highest
				q-value, divide .9 evenly among them */
			for(cur_edge = 0; cur_edge < total_edges; cur_edge++){
				float q_value = q_table[cur_state][cur_edge];
				if(!EDGE_IS_PLACED(cur_state, cur_edge) && q_value >= q_high){
					if(q_value == q_high){
						best = best | (int)pow(2, cur_edge);
						q_count++;
					}
					else{
						best = 0;
						best = best | (int)pow(2, cur_edge);
						q_count = 1;
					}
					q_high = q_value;
				}
			}
			placed = 0;
			/* This loop counts all the edges already placed and gives them 0% */
			for(cur_edge = 0; cur_edge < total_edges; cur_edge++){
				pi_table[cur_state][cur_edge] = 0;
				if(EDGE_IS_PLACED(cur_state, cur_edge))
					placed++;
			}
			low = .1 / (total_edges - placed - q_count);
			/* for the remaining unplaced edges, distribute percentages accordingly */
			the_chosen = ((total_edges - q_count) == placed) ? (1.0 / q_count)  : (.9 / q_count);
			for(cur_edge = 0; cur_edge < total_edges; cur_edge++){
				if(!EDGE_IS_PLACED(cur_state, cur_edge))
					pi_table[cur_state][cur_edge] = EDGE_IS_PLACED(best, cur_edge) ? the_chosen : low;
			}
		}
	}
}

//takes an edge index, and finds all the squares its
//adjacent too. returns the square index.
int edgeToSquare(int start, int edge_index){
	int x;
	for(x = start + 1; x < total_squares; x++){
		if(board[x].containsEdge(edge_index))
			return x;
	}
	return -1;
}

//takes a state, and the index of the edge placed.
//It then calculates if that placed edge completed 
//any squares.
int getPoints(int edge_index, int state){
	int square_index =-1;
	int potential_score = 0;	
	int edges_placed;	

	//get the first square the edge touches
	square_index = edgeToSquare(square_index, edge_index);
	edges_placed = board[square_index].calcFilled(state);
	if(edges_placed == 4)
		potential_score++;

	//check the second square, if there is one	
	square_index = edgeToSquare(square_index, edge_index);
  if(square_index != -1){
		edges_placed = board[square_index].calcFilled(state);
		if(edges_placed == 4)
			potential_score++;	
	}
	if(potential_score > 0 && !training)	
		cout << "Scored: " << potential_score << " point(s)\n";
	return potential_score;
}

/*NOTE: printTable is disabled if there are too many states. 2^22 is the cut off point */
void printTable(float **table){
	if(total_edges <= 22){
		for(cur_state = 0; cur_state < total_states; cur_state++){
			OUT_TYPE << "S" << cur_state << "\t";
			if(table[cur_state] == NULL) //initialize needed part of the table
				OUT_TYPE << "Computer never needed to learn this state\n";
			else{
				for(cur_edge = 0; cur_edge < total_edges; cur_edge++)
					OUT_TYPE << table[cur_state][cur_edge] << "\t";
				OUT_TYPE << endl;
			}
		}
		OUT_TYPE << endl;
	}
}
